
import pandas as pd
import json
import re
from pathlib import Path
import os
from datetime import datetime


def extract_numbers_from_text(text):
    """从文本中提取数字信息（年份、人口、面积）"""
    if pd.isna(text) or not isinstance(text, str):
        return {'year': None, 'population': None, 'area': None}

    text_clean = text.replace(' ', '').replace('\n', '').replace('\t', '')

    # 1. 提取年份
    year = None

    # 处理"前XXX年"格式（公元前）
    bc_year_match = re.search(r'前(\d{2,4})年', text_clean)
    if bc_year_match:
        try:
            year = -int(bc_year_match.group(1))  # 公元前用负数表示
        except:
            pass

    # 处理"公元XXX年"或"XXX年"格式（公元后）
    if not year:
        ad_year_match = re.search(r'公元(\d{3,4})年', text_clean)
        if ad_year_match:
            try:
                year = int(ad_year_match.group(1))
            except:
                pass

    # 处理直接的年份数字（优先4位数，其次3位数）
    if not year:
        year_match = re.search(r'(\d{4})年', text_clean)
        if not year_match:
            year_match = re.search(r'(\d{3})年', text_clean)

        if year_match:
            try:
                year_candidate = int(year_match.group(1))
                # 合理年份范围判断（公元前2000年到公元2000年）
                if year_candidate > 0 and year_candidate <= 2000:
                    year = year_candidate
            except:
                pass

    # 2. 提取人口数量
    population = None

    # 处理"XXX万户"格式
    population_wanhu_match = re.search(r'(\d+(?:\.\d+)?)万户', text_clean)
    if population_wanhu_match:
        try:
            population = int(float(population_wanhu_match.group(1)) * 10000)
        except:
            pass

    # 处理"XXX万口"或"XXX万人"格式
    if not population:
        population_wankou_match = re.search(r'(\d+(?:\.\d+)?)万[人口]', text_clean)
        if population_wankou_match:
            try:
                population = int(float(population_wankou_match.group(1)) * 10000)
            except:
                pass

    # 处理"XXX户"格式
    if not population:
        population_hu_match = re.search(r'(\d+)户', text_clean)
        if population_hu_match:
            try:
                # 按每户5人估算人口
                population = int(population_hu_match.group(1)) * 5
            except:
                pass

    # 处理"XXX口"格式
    if not population:
        population_kou_match = re.search(r'(\d+)口', text_clean)
        if population_kou_match:
            try:
                population = int(population_kou_match.group(1))
            except:
                pass

    # 处理直接的人口数字（如"俘人口三万余"）
    if not population:
        population_direct_match = re.search(r'人口(\d+)万', text_clean)
        if population_direct_match:
            try:
                population = int(population_direct_match.group(1)) * 10000
            except:
                pass

    # 3. 提取面积信息
    area = None  # 单位：平方千米

    # 处理"XXX里"（按1里=0.5千米计算，面积按正方形估算）
    area_li_match = re.search(r'(\d+)余里|(\d+)里', text_clean)
    if area_li_match:
        try:
            li_value = int(area_li_match.group(1) if area_li_match.group(1) else area_li_match.group(2))
            # 简单估算：如果是长度，按100里=50千米，面积按矩形100x50估算
            if li_value > 0:
                km_value = li_value * 0.5
                area = round(km_value * km_value / 10, 2)  # 简化估算
        except:
            pass

    # 处理"XXX亩"（1平方千米=1500亩）
    area_mu_match = re.search(r'(\d+)亩', text_clean)
    if area_mu_match and not area:
        try:
            area = round(int(area_mu_match.group(1)) / 1500, 2)
        except:
            pass

    # 处理"XXX平方里"
    area_squareli_match = re.search(r'(\d+)平方里', text_clean)
    if area_squareli_match and not area:
        try:
            # 1平方里=0.25平方千米
            area = round(int(area_squareli_match.group(1)) * 0.25, 2)
        except:
            pass

    return {'year': year, 'population': population, 'area': area}


def standardize_dynasty(dynasty):
    """标准化朝代名称"""
    if pd.isna(dynasty):
        return None

    dynasty_map = {
        '西周': '西周', '春秋战国': '春秋战国', '秦': '秦', '汉': '汉', '魏晋': '魏晋',
        '北魏': '北魏', '隋': '隋', '唐': '唐', '辽': '辽', '金': '金', '元': '元',
        '明': '明', '清': '清', '清代': '清', '明代': '明', '元代': '元', '金代': '金',
        '辽代': '辽', '唐代': '唐', '隋代': '隋', '汉代': '汉', '秦代': '秦'
    }

    dynasty_str = str(dynasty).strip()
    for key, value in dynasty_map.items():
        if key in dynasty_str:
            return value

    # 特殊处理
    if '先秦' in dynasty_str:
        return '先秦'
    elif '汉至唐' in dynasty_str:
        return '汉至唐'
    else:
        return dynasty_str


def extract_data_from_person_excel(file_path):
    """从人物Excel文件中提取人口、面积、年份数据"""
    try:
        # 获取所有sheet名称
        excel_file = pd.ExcelFile(file_path)
        all_sheets = excel_file.sheet_names
        print(f"发现 {len(all_sheets)} 个sheet: {all_sheets}")

        results = []
        sheet_stats = {}

        for sheet_name in all_sheets:
            print(f"\n正在处理sheet: {sheet_name}")
            df = pd.read_excel(file_path, sheet_name=sheet_name)

            # 检查必要的列是否存在
            required_columns = ['时期', '该人物在北京的活动']
            if not all(col in df.columns for col in required_columns):
                print(f"⚠️  {sheet_name} 缺少必要列，跳过")
                continue

            sheet_count = 0

            # 处理每一行数据
            for idx, row in df.iterrows():
                # 获取朝代（时期）
                dynasty = standardize_dynasty(row['时期'])
                # 获取活动描述
                activity_text = row['该人物在北京的活动']

                if dynasty and pd.notna(activity_text):
                    # 提取数据
                    extracted_data = extract_numbers_from_text(str(activity_text))

                    # 创建记录（至少有一个有效数据才保存）
                    if any([extracted_data['year'], extracted_data['population'], extracted_data['area']]):
                        record = {
                            'category': dynasty,
                            'year': extracted_data['year'],
                            'population': extracted_data['population'],
                            'area': extracted_data['area'],
                            'source_sheet': sheet_name,
                            'original_text': str(activity_text)[:200] + '...' if len(str(activity_text)) > 200 else str(
                                activity_text)
                        }
                        results.append(record)
                        sheet_count += 1

            sheet_stats[sheet_name] = sheet_count
            print(f"✅ 从 {sheet_name} 提取了 {sheet_count} 条有效记录")

        return results, sheet_stats

    except Exception as e:
        print(f"❌ 处理文件时出错: {str(e)}")
        return [], {}


def save_data_to_json(data, output_file):
    """保存数据到JSON文件"""
    try:
        # 创建输出目录
        output_dir = os.path.dirname(output_file)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 保存JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

        print(f"\n📊 数据已保存到: {output_file}")
        print(f"📈 总共提取了 {len(data)} 条有效记录")
        return True

    except Exception as e:
        print(f"❌ 保存文件时出错: {str(e)}")
        return False


def generate_summary_report(data, sheet_stats, output_file):
    """生成数据提取总结报告"""
    try:
        report_content = f"""# 人物数据中人口、面积、年份信息提取报告
生成时间: {datetime.now().strftime('%Y年%m月%d日 %H:%M:%S')}

## 1. 数据提取概况
- **总记录数**: {len(data)} 条
- **涉及sheet数**: {len(sheet_stats)} 个
- **数据字段**: 朝代(category)、年份(year)、人口(population)、面积(area)

## 2. 各sheet提取情况
| Sheet名称 | 提取记录数 |
|----------|------------|
"""

        for sheet_name, count in sheet_stats.items():
            report_content += f"| {sheet_name} | {count} |\n"

        # 统计各朝代数据
        dynasty_stats = {}
        year_count = 0
        population_count = 0
        area_count = 0

        for record in data:
            dynasty = record['category']
            if dynasty not in dynasty_stats:
                dynasty_stats[dynasty] = 0
            dynasty_stats[dynasty] += 1

            if record['year'] is not None:
                year_count += 1
            if record['population'] is not None:
                population_count += 1
            if record['area'] is not None:
                area_count += 1

        report_content += f"""
## 3. 各朝代数据分布
| 朝代 | 记录数 |
|------|--------|
"""

        for dynasty, count in sorted(dynasty_stats.items()):
            report_content += f"| {dynasty} | {count} |\n"

        report_content += f"""
## 4. 数据类型统计
- **包含年份信息**: {year_count} 条 ({round(year_count / len(data) * 100, 1)}%)
- **包含人口信息**: {population_count} 条 ({round(population_count / len(data) * 100, 1)}%)
- **包含面积信息**: {area_count} 条 ({round(area_count / len(data) * 100, 1)}%)

## 5. 数据说明
1. **年份格式**: 公元前年份用负数表示（如前221年表示为-221），公元后年份用正数表示
2. **人口单位**: 统一为"人"（万户、万口已转换为具体人数）
3. **面积单位**: 统一为"平方千米"（里、亩等已进行单位转换）
4. **估算说明**: 部分数据基于历史记载进行合理估算（如户数按每户5人估算）

## 6. 数据质量说明
- 提取的数据来自"该人物在北京的活动"字段中的历史记载
- 部分古代数据可能存在估算成分，仅供研究参考
- 建议结合其他历史文献进行交叉验证
"""

        # 保存报告
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(report_content)

        print(f"📄 报告已保存到: {output_file}")
        return True

    except Exception as e:
        print(f"❌ 生成报告时出错: {str(e)}")
        return False


def main():
    """主函数"""
    print("=" * 60)
    print("        人物数据人口、面积、年份信息提取工具")
    print("=" * 60)

    # 配置文件路径
    input_file = './data/19人物 - 总数据和各朝代数据.xlsx'
    output_dir = './outputs'
    json_output_file = os.path.join(output_dir, 'person_population_area_data.json')
    report_output_file = os.path.join(output_dir, 'data_extraction_report.md')

    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"📁 创建输出目录: {output_dir}")

    # 1. 提取数据
    print(f"\n🔍 开始从Excel文件提取数据: {input_file}")
    extracted_data, sheet_stats = extract_data_from_person_excel(input_file)

    if not extracted_data:
        print("❌ 未提取到任何有效数据")
        return

    # 2. 保存JSON数据
    save_data_to_json(extracted_data, json_output_file)

    # 3. 生成报告
    generate_summary_report(extracted_data, sheet_stats, report_output_file)

    # 4. 显示示例数据
    print(f"\n📋 前5条数据示例:")
    print("-" * 80)
    for i, record in enumerate(extracted_data[:5], 1):
        print(f"记录 {i}:")
        print(f"  朝代: {record['category']}")
        print(f"  年份: {record['year']} (公元前为负数)" if record['year'] is not None else "  年份: 无")
        print(f"  人口: {record['population']} 人" if record['population'] is not None else "  人口: 无")
        print(f"  面积: {record['area']} 平方千米" if record['area'] is not None else "  面积: 无")
        print(f"  来源: {record['source_sheet']}")
        print(f"  原文: {record['original_text'][:100]}...")
        print()

    print("=" * 60)
    print("✅ 数据提取任务完成！")
    print(f"📊 输出文件位置: {output_dir}")
    print("=" * 60)


if __name__ == '__main__':
    main()